import os
import torch
from ptflops import get_model_complexity_info
from torch.utils.data import DataLoader

from seg_dataset import mydata
from train_utils.train_and_eval import evaluate
from I_FCN import VGG16



def create_model(aux, num_classes, pretrain=False):

    model = VGG16().to('cuda')

    flops, params = get_model_complexity_info(model, (3, 384, 384), as_strings=True, print_per_layer_stat=True)
    print('FLOPs:{}'.format(flops))
    print('Params:{}'.format(params))
    model.load_state_dict(torch.load('F:\\lunwen\\save_weights\\vgg_decoder_aspp_unit_384_nonormal8.pth')['model'])

    return model


def main(args):
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    batch_size = args.batch_size
    # segmentation nun_classes + background
    num_classes = args.num_classes + 1

    # 用来保存训练以及验证过程中信息


    iss =384

    val_data = mydata('D:\\study\\pytorch_study\\seg_thryoid_picture\\datasets', iss, 'val')


    val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False, drop_last=True)

    model = create_model(aux=args.aux, num_classes=num_classes)
    confmat = evaluate(model, val_loader, device=device, num_classes=num_classes)
    val_info = str(confmat)
    print(val_info)



def parse_args():
    import argparse
    parser = argparse.ArgumentParser(description="pytorch deeplabv3 training")

    parser.add_argument("--data-path", default="/data/", help="VOCdevkit root")
    parser.add_argument("--num-classes", default=1, type=int)
    parser.add_argument("--aux", default=False, type=bool, help="auxilier loss")#me
    parser.add_argument("--device", default="cuda", help="training device")
    parser.add_argument("-b", "--batch-size", default=6, type=int)#me
    parser.add_argument("--epochs", default=20, type=int, metavar="N",
                        help="number of total epochs to train")#me

    parser.add_argument('--lr', default=0.0001, type=float, help='initial learning rate')
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                        help='momentum')
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                        metavar='W', help='weight decay (default: 1e-4)',
                        dest='weight_decay')#me
    parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
    parser.add_argument('--resume', default='F:\\lunwen\\\save_weights\\resnet_backbone1.pth', help='resume from checkpoint')
    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    # Mixed precision training parameters
    parser.add_argument("--amp", default=False, type=bool,
                        help="Use torch.cuda.amp for mixed precision training")

    args = parser.parse_args()

    return args


if __name__ == '__main__':
    args = parse_args()

    if not os.path.exists("./save_weights"):
        os.mkdir("./save_weights")

    main(args)

